/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef __AICPU_RPC_SERVERV2_H__
#define __AICPU_RPC_SERVERV2_H__

#include "common/aicpu_hccl_def.h"
#include "stream_pub.h"

class AicpuRpcServerV2 {
public:
    AicpuRpcServerV2() = default;
    ~AicpuRpcServerV2() = default;
    HcclResult Init(const HcclMC2WorkSpace &workspaceInfo);
    void Reset();
    uint8_t GetTaskType() const;
    HcclMsgArea* GetHcclMsgArea(void);
    HcclMsg *GetMsgWorkSpace(void);
    uint64_t GetWorkSpaceCommitArea(void);
    uint64_t GetWorkSpaceFinishArea(void);
    uint64_t GetFinishAddr(int32_t idx);
    uint64_t GetCommitareaAddr(int32_t idx);
    uint64_t GetFinishAddrByHandleId(HcclHandle handleId);
    bool GetIsFinalize(void) const;
    HcclResult AddCcoreNotify(uint16_t streamId, u64 recordAddr, uint32_t turnNum, hccl::Stream *stream);
    HcclResult AddCcoreWait(uint16_t streamId, u64 waitAddr, uint32_t turnNum, hccl::Stream *stream, bool isLast);
    HcclResult ResetCommitTaskAdd(u16 streamId, hccl::Stream *stream);
    void WriteFinishWhenAllFinalize();
    void WriteRestartFlag();
    uint32_t GetMsgPos(void) const;
    void SetMsgRepeatCnt(u8 repeatCnt);
    int32_t GetMsgRepeatCnt(HcclHandle handleId);
    int32_t GetMsgHandlePos(HcclHandle handleId);
    void SetMsgPos(uint32_t pos);
    void SetIsFinalize(bool finalize);
    void PrintAllHcclMsgArea(u32 rankSize);
    void SetMsgPosForKernel(uint32_t msgPos);
    uint32_t GetMsgPosForKernel(void);
    void SetMsgHandlePos(uint32_t msgPos, HcclHandle handleId);
    void SetNeedRetryFlag(bool needRetryFlag);
    bool ReadAddrMsg(HcclMsg *hcclMsg, HcclMsg *msgList, uint32_t msgPos, u32 rankSize);
    static bool ReadAddrMsg(AicpuRpcServerV2 *rpcInst, HcclMsg *hcclMsg, HcclMsg *msgList, uint32_t msgPos, u32 rankSize);
    HcclMsgExt* GetHcclMsgExtPtr();
    std::string HcclMsgToStr(HcclMsgV1 *msg);
    std::string HcclMsgToTtr(HcclMsg *msg);
    HcclResult ProcessExpectPrepareMsg(CommonHcclMsg *hcclMsg, uint8_t expectId);

public:
    void SetStepSize(u8 stepSize) { curStepSize_ = stepSize; };
    u8 GetStepSize() const { return curStepSize_; }
    u64 GetTurnNumAddr() { return turnNumAddr_; }

private:
    void HcclMsgExtToStr(const HcclMsgExt &msg, u32 rankSize);
    bool ReadValidMsg(HcclMsg *rMsg, HcclMsg *msg, bool reset, bool needReProcess,
        uint32_t msgPos, u32 rankSize);
    uint64_t GenXorForMsgExt(int32_t idx, u32 rankSize);
    bool ReadValidMsgExtArea(int32_t idx, u32 rankSize);
    void GenMsgByTaskParam(AivAicpuOpParam *outMsg);

private:
    uint64_t workSpace_ = 0;
    HcclMsgArea *hcclMsgArea_= nullptr;
    uint32_t repeatCnt_[AC_MSG_CNT];
    int8_t handleIdMsgPosition_[MC2_HCCL_MAX_HANDLE_ID];
    uint64_t streamId_;
    uint32_t msgPosForKernel_;
    uint32_t msgPos_;
    bool needReProcess_ = false;
    bool isFinalize_;
    std::shared_ptr<HcclMsgExt> msgExt_ = std::make_shared<HcclMsgExt>();
    u64 prepareTime_;                                                           // 记录 Prepare 消息的时间
    u8 eventPrintTurn_;                                                         // 记录打印event的turn
    u8 curStepSize_ = 0U;
    u64 turnNumAddr_;
};

#endif  // __AICPU_RPC_SERVERV2_H__